/*
 * wifi_test.c
 * 
 * 使用hostapd启动WiFi接入点或使用wpa_supplicant连接WiFi的测试程序
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <signal.h>
#include "wifi_ap.h"
#include "wifi_sta.h"

// 运行模式枚举
typedef enum {
    MODE_AP = 0,
    MODE_STA = 1
} run_mode_t;

// 全局状态变量
wifi_ap_state_t g_ap_state = {0};
wifi_sta_state_t g_sta_state = {0};
int g_running = 1;
run_mode_t g_mode = MODE_AP;

/**
 * 信号处理函数
 */
void signal_handler(int sig)
{
    if (sig == SIGINT || sig == SIGTERM) {
        printf("\nReceived termination signal, shutting down...\n");
        g_running = 0;
    }
}

/**
 * 显示程序帮助信息
 */
void show_help(void)
{
    printf("\nWiFi Test Program\n");
    printf("Usage: sudo ./wifi_test [options]\n");
    printf("Options:\n");
    printf("  -o <mode>        Operation mode (ap/sta, default: ap)\n");
    printf("  -i <interface>   WiFi interface name (default: wlan0)\n");
    printf("  -s <ssid>        WiFi SSID\n");
    printf("  -p <password>    WiFi password (at least 8 characters for AP mode)\n");
    printf("  -c <channel>     WiFi channel (1-13, default: 6, AP mode only)\n");
    printf("  -m <mode>        WiFi mode (g/n/ac, default: g, AP mode only)\n");
    printf("  -h               Show this help message\n");
}

/**
 * 打印WiFi AP状态
 */
void print_ap_status(const wifi_ap_state_t *state)
{
    printf("\n========== WiFi AP Status ==========\n");
    printf("State: %s\n", state->running ? "Running" : "Stopped");
    
    if (state->running) {
        printf("Interface: %s\n", state->config.interface);
        printf("SSID: %s\n", state->config.ssid);
        printf("Channel: %s\n", state->config.channel);
        printf("Mode: %s\n", state->config.hw_mode);
        printf("Encryption: WPA-%s\n", state->config.wpa);
        printf("PID: %d\n", state->pid);
    }
    printf("====================================\n");
}

/**
 * 打印WiFi STA状态
 */
void print_sta_status(const wifi_sta_state_t *state)
{
    printf("\n========== WiFi STA Status ==========\n");
    printf("State: %s\n", state->connected ? "Connected" : "Disconnected");
    
    if (state->connected) {
        printf("Interface: %s\n", state->config.interface);
        printf("Connected to SSID: %s\n", state->config.ssid);
        printf("IP Address: %s\n", strlen(state->ip_address) > 0 ? state->ip_address : "Not assigned");
        printf("PID: %d\n", state->pid);
    }
    printf("====================================\n");
}

/**
 * 使用udhcpd配置DHCP服务
 */
wifi_ap_error_t configure_dhcp(const char *interface, const char *ip_range_start, const char *ip_range_end, const char *netmask)
{
    char cmd[512];
    FILE *fp;
    const char *udhcpd_conf_path = "/tmp/udhcpd.conf";
    const char *pid_file = "/var/run/udhcpd.pid";
    wifi_ap_error_t ret;

    // 停止已运行的udhcpd
    if (file_exists(pid_file)) {
        execute_command("killall udhcpd", "Warning: Failed to kill existing udhcpd");
        unlink(pid_file);
    }

    // 创建udhcpd配置文件
    fp = fopen(udhcpd_conf_path, "w");
    if (fp == NULL) {
        perror("Failed to create udhcpd config file");
        return WIFI_AP_ERROR_CONFIG;
    }

    fprintf(fp, "start %s\n", ip_range_start);
    fprintf(fp, "end %s\n", ip_range_end);
    fprintf(fp, "interface %s\n", interface);
    fprintf(fp, "option subnet %s\n", netmask);
    fprintf(fp, "option router %s\n", ip_range_start);
    fprintf(fp, "option dns 8.8.8.8\n");
    fprintf(fp, "option lease 86400\n");
    fprintf(fp, "pidfile %s\n", pid_file);

    fclose(fp);

    // 启动udhcpd
    sprintf(cmd, "udhcpd %s", udhcpd_conf_path);
    ret = execute_command(cmd, "Failed to start udhcpd");
    if (ret != WIFI_AP_SUCCESS) {
        return WIFI_AP_ERROR_PROCESS;
    }

    return WIFI_AP_SUCCESS;
}

int main(int argc, char *argv[])
{
    wifi_ap_config_t ap_config;
    wifi_sta_config_t sta_config;
    wifi_ap_error_t ret;
    int opt;

    // 检查是否以root权限运行
    if (geteuid() != 0) {
        fprintf(stderr, "Error: This program must be run as root\n");
        show_help();
        return 1;
    }

    // 初始化AP模式配置为默认值
    ret = wifi_ap_init_config(&ap_config);
    if (ret != WIFI_AP_SUCCESS) {
        fprintf(stderr, "Failed to initialize AP config: %s\n", 
                wifi_ap_get_error_string(ret));
        return 1;
    }

    // 初始化STA模式配置为默认值
    ret = wifi_sta_init_config(&sta_config);
    if (ret != WIFI_AP_SUCCESS) {
        fprintf(stderr, "Failed to initialize STA config: %s\n", 
                wifi_ap_get_error_string(ret));
        return 1;
    }

    // 解析命令行参数
    while ((opt = getopt(argc, argv, "o:i:s:p:c:m:h")) != -1) {
        switch (opt) {
            case 'o':
                if (strcmp(optarg, "sta") == 0) {
                    g_mode = MODE_STA;
                } else if (strcmp(optarg, "ap") != 0) {
                    fprintf(stderr, "Error: Invalid mode. Use 'ap' or 'sta'\n");
                    show_help();
                    return 1;
                }
                break;
            case 'i':
                strncpy(ap_config.interface, optarg, sizeof(ap_config.interface) - 1);
                strncpy(sta_config.interface, optarg, sizeof(sta_config.interface) - 1);
                break;
            case 's':
                strncpy(ap_config.ssid, optarg, sizeof(ap_config.ssid) - 1);
                strncpy(sta_config.ssid, optarg, sizeof(sta_config.ssid) - 1);
                break;
            case 'p':
                if (g_mode == MODE_AP && strlen(optarg) < 8) {
                    fprintf(stderr, "Error: Password must be at least 8 characters for AP mode\n");
                    show_help();
                    return 1;
                }
                strncpy(ap_config.password, optarg, sizeof(ap_config.password) - 1);
                strncpy(sta_config.password, optarg, sizeof(sta_config.password) - 1);
                break;
            case 'c':
                strncpy(ap_config.channel, optarg, sizeof(ap_config.channel) - 1);
                break;
            case 'm':
                strncpy(ap_config.hw_mode, optarg, sizeof(ap_config.hw_mode) - 1);
                break;
            case 'h':
                show_help();
                return 0;
            default:
                show_help();
                return 1;
        }
    }

    // 设置信号处理
    signal(SIGINT, signal_handler);
    signal(SIGTERM, signal_handler);

    if (g_mode == MODE_AP) {
        printf("WiFi AP Test Program\n");
        printf("Configuring WiFi AP with the following settings:\n");
        printf("  Interface: %s\n", ap_config.interface);
        printf("  SSID: %s\n", ap_config.ssid);
        printf("  Password: ********\n");
        printf("  Channel: %s\n", ap_config.channel);
        printf("  Mode: %s\n", ap_config.hw_mode);

        // 设置WiFi接口IP地址
        printf("\nSetting up network interface...\n");
        ret = wifi_ap_set_ip(ap_config.interface, "192.168.42.1", "255.255.255.0");
        if (ret != WIFI_AP_SUCCESS) {
            fprintf(stderr, "Failed to set IP address: %s\n", 
                    wifi_ap_get_error_string(ret));
            return 1;
        }

        // 配置DHCP服务
        printf("Configuring DHCP server...\n");
        ret = configure_dhcp(ap_config.interface, "192.168.42.100", "192.168.42.200", "255.255.255.0");
        if (ret != WIFI_AP_SUCCESS) {
            fprintf(stderr, "Warning: Failed to configure DHCP server: %s\n", 
                    wifi_ap_get_error_string(ret));
            // 继续执行，即使DHCP配置失败
        }

        // 启动WiFi AP
        printf("Starting WiFi AP...\n");
        ret = wifi_ap_start(&ap_config, &g_ap_state);
        if (ret != WIFI_AP_SUCCESS) {
            fprintf(stderr, "Failed to start WiFi AP: %s\n", 
                    wifi_ap_get_error_string(ret));
            
            // 清理资源
            wifi_ap_cleanup(&ap_config);
            return 1;
        }

        printf("WiFi AP started successfully!\n");
        print_ap_status(&g_ap_state);
        printf("\nPress Ctrl+C to stop...\n");

        // 主循环
        while (g_running) {
            // 每秒检查一次状态
            sleep(1);
            
            // 检查AP状态
            ret = wifi_ap_check_status(&g_ap_state);
            if (ret != WIFI_AP_SUCCESS && g_running) {
                fprintf(stderr, "WiFi AP stopped unexpectedly: %s\n", 
                        wifi_ap_get_error_string(ret));
                g_running = 0;
            }
        }

        // 停止WiFi AP
        printf("\nStopping WiFi AP...\n");
        ret = wifi_ap_stop(&g_ap_state);
        if (ret != WIFI_AP_SUCCESS) {
            fprintf(stderr, "Failed to stop WiFi AP cleanly: %s\n", 
                    wifi_ap_get_error_string(ret));
        }

        // 清理资源
        printf("Cleaning up resources...\n");
        wifi_ap_cleanup(&ap_config);
    } else {
        // STA模式
        printf("WiFi STA Test Program\n");
        printf("Configuring WiFi STA with the following settings:\n");
        printf("  Interface: %s\n", sta_config.interface);
        printf("  SSID: %s\n", sta_config.ssid);
        printf("  Password: %s\n", strlen(sta_config.password) > 0 ? "********" : "(none)");

        // 检查SSID是否为空
        if (strlen(sta_config.ssid) == 0) {
            fprintf(stderr, "Error: SSID cannot be empty in STA mode\n");
            show_help();
            return 1;
        }

        // 启动WiFi STA
        printf("\nStarting WiFi STA and connecting to %s...\n", sta_config.ssid);
        ret = wifi_sta_start(&sta_config, &g_sta_state);
        if (ret != WIFI_AP_SUCCESS) {
            fprintf(stderr, "Failed to start WiFi STA: %s\n", 
                    wifi_ap_get_error_string(ret));
            
            // 清理资源
            wifi_sta_cleanup(&sta_config);
            return 1;
        }

        printf("WiFi STA started successfully!\n");
        print_sta_status(&g_sta_state);
        printf("\nPress Ctrl+C to stop...\n");

        // 主循环
        while (g_running) {
            // 每秒检查一次状态
            sleep(1);
            
            // 检查STA状态
            ret = wifi_sta_check_status(&g_sta_state);
            if (ret != WIFI_AP_SUCCESS && g_running) {
                fprintf(stderr, "WiFi STA connection issue: %s\n", 
                        wifi_ap_get_error_string(ret));
                g_running = 0;
            }
        }

        // 停止WiFi STA
        printf("\nStopping WiFi STA...\n");
        ret = wifi_sta_stop(&g_sta_state);
        if (ret != WIFI_AP_SUCCESS) {
            fprintf(stderr, "Failed to stop WiFi STA cleanly: %s\n", 
                    wifi_ap_get_error_string(ret));
        }

        // 清理资源
        printf("Cleaning up resources...\n");
        wifi_sta_cleanup(&sta_config);
    }

    printf("Program exited successfully.\n");
    return 0;
}